# factor_builder.py
import pandas as pd
import numpy as np


def calculate_momentum_factor(data_dict, lookback_period=20):
    """
    为每只股票计算动量因子：过去N天的收益率
    """
    factor_values = {}

    for ts_code, df in data_dict.items():
        # 计算过去lookback_period天的收益率
        df['close'] = df['close'].astype(float)
        df['momentum'] = df['close'].pct_change(periods=lookback_period)
        # 只保留需要的列
        factor_series = df['momentum']
        factor_values[ts_code] = factor_series

    # 转换为DataFrame，索引为日期，列为股票代码
    factor_df = pd.DataFrame(factor_values)
    return factor_df


def generate_factor_signals(factor_df, top_n=10):
    """
    根据因子值生成交易信号：做多因子值最高的top_n只股票
    """
    signals = {}

    for date in factor_df.index:
        daily_factors = factor_df.loc[date]
        # 剔除NaN值
        valid_factors = daily_factors.dropna()
        if len(valid_factors) < top_n:
            continue

        # 找到因子值最高的top_n只股票
        top_stocks = valid_factors.nlargest(top_n).index.tolist()
        signals[date] = top_stocks

    return signals


if __name__ == "__main__":
    # 加载数据
    import pickle

    with open('stock_data.pkl', 'rb') as f:
        data = pickle.load(f)

    # 计算动量因子
    factor_df = calculate_momentum_factor(data)

    # 生成交易信号
    signals = generate_factor_signals(factor_df)

    # 保存信号
    with open('trading_signals.pkl', 'wb') as f:
        pickle.dump(signals, f)
